import pandas as pd
import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq


def read_embeddings_from_parquet(file_path, index_col='image_id'):
    """
    从Parquet文件中读取数据，并将其转换为以指定列为键的字典。
    """
    try:
        print(f"正在从 {file_path} 读取数据...")
        df = pd.read_parquet(file_path)
        print("数据读取成功！")

        if index_col not in df.columns:
            print(f"错误: 索引列 '{index_col}' 在文件中未找到。")
            return None
        
        # 将索引列转换为字符串，以确保字典键的类型统一
        df[index_col] = df[index_col].astype(str)
        
        print(f"正在将数据转换为以 '{index_col}' 为键的字典...")
        df.set_index(index_col, inplace=True)
        embeddings_dict = df.to_dict(orient='index')
        print("转换完成！")
        
        return embeddings_dict
    except FileNotFoundError:
        print(f"错误：文件未找到于路径 '{file_path}'")
        return None
    except Exception as e:
        print(f"处理文件时发生未知错误: {e}")
        return None


def save_embeddings_to_parquet(ans, output_path):
    """
    将包含嵌入向量的字典保存为Parquet文件。

    Args:
        ans (dict): 待保存的字典，键为 image_id。
        output_path (str): Parquet文件的输出路径。
    """
    if not ans:
        print("警告: 嵌入向量字典为空，没有内容可保存。")
        return
    for key, item in ans.items():
        item['image_id'] = key
    # 将字典的值转换为DataFrame
    df = pd.DataFrame.from_dict(ans, orient='index')
    
    # 将DataFrame转换为PyArrow Table
    try:
        
        table = pa.Table.from_pandas(df)
        pq.write_table(table, output_path)
        print(f"嵌入向量已成功保存至: {output_path}")
    except Exception as e:
        print(f"错误: 保存 Parquet 文件失败。 {e}")


def merge_qv2R_parquet_files(q_path,v_path,output_path):
    q = read_embeddings_from_parquet(q_path)
    v = read_embeddings_from_parquet(v_path)
    R = q | v
    save_embeddings_to_parquet(R,output_path)

def merge_oR2I_parquet_files(o_path,r_path,output_path):
    o = read_embeddings_from_parquet(o_path)
    r = read_embeddings_from_parquet(r_path)
    I = o | r
    save_embeddings_to_parquet(I,output_path)


if __name__ == "__main__":
    q_path = 'outputs/embeddings/no_train/baseline_embeddings_q_norm.parquet'
    v_path = 'outputs/embeddings/no_train/baseline_embeddings_v_norm.parquet'
    output_path = 'outputs/embeddings/no_train/baseline_embeddings_R_norm.parquet'
    merge_qv2R_parquet_files(q_path,v_path,output_path)

    o_path = 'outputs/embeddings/no_train/baseline_embeddings_o_norm.parquet'
    r_path = 'outputs/embeddings/no_train/baseline_embeddings_R_norm.parquet'
    output_path = 'outputs/embeddings/no_train/baseline_embeddings_I_norm.parquet'
    merge_oR2I_parquet_files(o_path,r_path,output_path)
    